#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Minimal web interface to cve-search to display the last entries
# and view a specific CVE.
#
# Software is free software released under the "GNU Affero General Public License v3.0"
#
# Copyright (c) 2013-2018  Alexandre Dulaunoy - a@foo.be
# Copyright (c) 2014-2018  Pieter-Jan Moreels - pieterjan.moreels@gmail.com

import os
import re
import subprocess
import sys
import urllib
from io import TextIOWrapper, BytesIO

import jinja2
from flask import abort, jsonify, request, redirect, render_template, send_file
from flask_login import (
    LoginManager,
    current_user,
    login_user,
    logout_user,
    login_required,
)
from redis import exceptions as redisExceptions

_runPath = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(_runPath, ".."))

from lib.CVEs import CveHandler
from lib.DatabaseLayer import (
    getRules,
    getCVEs,
    getDBStats,
    getToken,
    generateToken,
    changePassword,
    getWhitelist,
    getBlacklist,
)
from sbin.db_blacklist import (
    insertBlacklist,
    countBlacklist,
    dropBlacklist,
    importBlacklist,
    exportBlacklist,
    removeBlacklist,
    updateBlacklist,
)
from sbin.db_whitelist import (
    insertWhitelist,
    countWhitelist,
    dropWhitelist,
    importWhitelist,
    exportWhitelist,
    removeWhitelist,
    updateWhitelist,
)
from lib.Query import getVersionsOfProduct, getBrowseList
from lib.Toolkit import tk_compile
from lib.Authentication import AuthenticationHandler
from lib.Config import Configuration
from lib.PluginManager import PluginManager
from lib.User import User
from web.minimal import Minimal
from web.advanced_api import Advanced_API


class Index(Minimal, Advanced_API):
    #############
    # Variables #
    #############

    def __init__(self):
        # TODO: make auth handler and plugin manager singletons
        Advanced_API.__init__(self)
        Minimal.__init__(self)
        self.minimal = False
        self.auth_handler = AuthenticationHandler()
        self.plugManager = PluginManager()
        self.login_manager = LoginManager()
        self.plugManager.loadPlugins()
        self.login_manager.init_app(self.app)
        self.login_manager.user_loader(self.load_user)
        self.redisdb = Configuration.getRedisVendorConnection()

        self.defaultFilters.update(
            {
                "blacklistSelect": "on",
                "whitelistSelect": "on",
                "unlistedSelect": "show",
            }
        )
        self.args.update({"minimal": False})
        self.pluginArgs = {
            "current_user": current_user,
            "plugin_manager": self.plugManager,
        }

        routes = [
            {"r": "/cve/<cveid>", "m": ["GET"], "f": self.cve},
            {"r": "/_get_plugins", "m": ["GET"], "f": self._get_plugins},
            {"r": "/plugin/_get_cve_actions", "m": ["GET"], "f": self._get_cve_actions},
            {"r": "/plugin/<plugin>", "m": ["GET"], "f": self.openPlugin},
            {
                "r": "/plugin/<plugin>/subpage/<page>",
                "m": ["GET"],
                "f": self.openPluginSubpage,
            },
            {
                "r": "/plugin/<plugin>/_cve_action/<action>",
                "m": ["GET"],
                "f": self._jsonCVEAction,
            },
            {"r": "/login", "m": ["POST"], "f": self.login_check},
            {"r": "/logout", "m": ["GET"], "f": self.logout},
            {"r": "/admin", "m": ["GET"], "f": self.admin},
            {"r": "/admin/", "m": ["GET"], "f": self.admin},
            {"r": "/admin/change_pass", "m": ["GET"], "f": self.change_pass},
            {"r": "/admin/request_token", "m": ["GET"], "f": self.request_token},
            {"r": "/admin/updatedb", "m": ["GET"], "f": self.updatedb},
            {"r": "/admin/whitelist/import", "m": ["POST"], "f": self.listImport},
            {"r": "/admin/blacklist/import", "m": ["POST"], "f": self.listImport},
            {"r": "/admin/whitelist/export", "m": ["GET"], "f": self.listExport},
            {"r": "/admin/blacklist/export", "m": ["GET"], "f": self.listExport},
            {"r": "/admin/whitelist/drop", "m": ["POST"], "f": self.listDrop},
            {"r": "/admin/blacklist/drop", "m": ["POST"], "f": self.listDrop},
            {"r": "/admin/whitelist", "m": ["GET"], "f": self.listView},
            {"r": "/admin/blacklist", "m": ["GET"], "f": self.listView},
            {"r": "/admin/addToList", "m": ["GET"], "f": self.listAdd},
            {"r": "/admin/removeFromList", "m": ["GET"], "f": self.listRemove},
            {"r": "/admin/editInList", "m": ["GET"], "f": self.listEdit},
            {"r": "/admin/listmanagement", "m": ["GET"], "f": self.listManagement},
            {
                "r": "/admin/listmanagement/<vendor>",
                "m": ["GET"],
                "f": self.listManagement,
            },
            {
                "r": "/admin/listmanagement/<vendor>/<product>",
                "m": ["GET"],
                "f": self.listManagement,
            },
            {
                "r": "/admin/listmanagement/add",
                "m": ["GET"],
                "f": self.listManagementAdd,
            },
            {"r": "/login", "m": ["POST"], "f": self.login_check},
        ]
        for route in routes:
            self.addRoute(route)

    #############
    # Functions #
    #############
    def generate_full_query(self, f):
        query = self.generate_minimal_query(f)
        if current_user.is_authenticated():
            if f["blacklistSelect"] == "on":
                regexes = getRules("blacklist")
                if len(regexes) != 0:
                    exp = "^(?!" + "|".join(regexes) + ")"
                    query.append(
                        {
                            "$or": [
                                {"vulnerable_configuration": re.compile(exp)},
                                {"vulnerable_configuration": {"$exists": False}},
                                {"vulnerable_configuration": []},
                            ]
                        }
                    )
            if f["whitelistSelect"] == "hide":
                regexes = getRules("whitelist")
                if len(regexes) != 0:
                    exp = "^(?!" + "|".join(regexes) + ")"
                    query.append(
                        {
                            "$or": [
                                {"vulnerable_configuration": re.compile(exp)},
                                {"vulnerable_configuration": {"$exists": False}},
                                {"vulnerable_configuration": []},
                            ]
                        }
                    )
            if f["unlistedSelect"] == "hide":
                wlregexes = tk_compile(getRules("whitelist"))
                blregexes = tk_compile(getRules("blacklist"))
                query.append(
                    {
                        "$or": [
                            {"vulnerable_configuration": {"$in": wlregexes}},
                            {"vulnerable_configuration": {"$in": blregexes}},
                        ]
                    }
                )
        return query

    def markCPEs(self, cve):
        blacklist = tk_compile(getRules("blacklist"))
        whitelist = tk_compile(getRules("whitelist"))

        for conf in cve["vulnerable_configuration"]:
            conf["list"] = "none"
            conf["match"] = "none"
            for w in whitelist:
                if w.match(conf["id"]):
                    conf["list"] = "white"
                    conf["match"] = w
            for b in blacklist:
                if b.match(conf["id"]):
                    conf["list"] = "black"
                    conf["match"] = b
        return cve

    def filter_logic(self, filters, skip, limit=None):
        query = self.generate_full_query(filters)
        limit = limit if limit else self.args["pageLength"]
        cve = getCVEs(limit=limit, skip=skip, query=query)
        # marking relevant records
        if current_user.is_authenticated():
            if filters["whitelistSelect"] == "on":
                cve["results"] = self.list_mark("white", cve["results"])
            if filters["blacklistSelect"] == "mark":
                cve["results"] = self.list_mark("black", cve["results"])
        self.plugManager.mark(cve, **self.pluginArgs)
        return cve

    def addCPEToList(self, cpe, listType, cpeType=None):
        def addCPE(cpe, cpeType, funct):
            return True if funct(cpe, cpeType) else False

        if not cpeType:
            cpeType = "cpe"

        if listType.lower() in ("blacklist", "black", "b", "bl"):
            return addCPE(cpe, cpeType, insertBlacklist)
        if listType.lower() in ("whitelist", "white", "w", "wl"):
            return addCPE(cpe, cpeType, insertWhitelist)

    def list_mark(self, listed, cveList):
        if listed not in ["white", "black"]:
            return list(cveList)
        items = tk_compile(getRules(listed + "list"))
        # check the cpes (full or partially) in the black/whitelist
        for i, cve in enumerate(
            list(cveList)
        ):  # the list() is to ensure we don't have a pymongo cursor object
            for c in cve["vulnerable_configuration"]:
                if any(regex.match(c) for regex in items):
                    cveList[i][listed + "listed"] = "yes"
        return cveList

    def filterUpdateField(self, data):
        if not data:
            return data
        returnvalue = []
        for line in data.split("\n"):
            if (
                not line.startswith("[+]Success to create index")
                and not line == "Not modified"
                and not line.startswith("Starting")
            ):
                returnvalue.append(line)
        return "\n".join(returnvalue)

    def adminInfo(self, output=None):
        return {
            "stats": getDBStats(True),
            "plugins": self.plugManager.getPlugins(),
            "updateOutput": self.filterUpdateField(output),
            "token": getToken(current_user.id),
        }

    # user management
    def load_user(self, id):
        return User.get(id, self.auth_handler)

    ##########
    # ROUTES #
    ##########
    # /cve/<cveid>
    def cve(self, cveid):
        cveid = cveid.upper()
        cvesp = CveHandler(
            rankinglookup=True,
            namelookup=True,
            via4lookup=True,
            capeclookup=True,
            subscorelookup=True,
        )
        cve = cvesp.getcve(cveid=cveid)
        if cve is None:
            return render_template(
                "error.html", status={"except": "cve-not-found", "info": {"cve": cveid}}
            )
        cve = self.markCPEs(cve)

        self.plugManager.onCVEOpen(cveid, **self.pluginArgs)
        pluginData = self.plugManager.cvePluginInfo(cveid, **self.pluginArgs)
        return render_template("cve.html", cve=cve, plugins=pluginData)

    # /_get_plugins
    def _get_plugins(self):
        if (
            not current_user.is_authenticated()
        ):  # Don't show plugins requiring auth if not authenticated
            plugins = [
                {"name": x.getName(), "link": x.getUID()}
                for x in self.plugManager.getWebPluginsWithPage(**self.pluginArgs)
                if not x.requiresAuth
            ]
        else:
            plugins = [
                {"name": x.getName(), "link": x.getUID()}
                for x in self.plugManager.getWebPluginsWithPage(**self.pluginArgs)
            ]
        return jsonify({"plugins": plugins})

    # /plugin/_get_cve_actions
    def _get_cve_actions(self):
        cve = request.args.get("cve", type=str)
        if (
            not current_user.is_authenticated()
        ):  # Don't show actions requiring auth if not authenticated
            actions = [
                x
                for x in self.plugManager.getCVEActions(cve, **self.pluginArgs)
                if not x["auth"]
            ]
        else:
            actions = self.plugManager.getCVEActions(cve, **self.pluginArgs)
        return jsonify({"actions": actions})

    # /plugin/<plugin>
    def openPlugin(self, plugin):
        if (
            self.plugManager.requiresAuth(plugin)
            and not current_user.is_authenticated()
        ):
            return render_template("requiresAuth.html")
        else:
            page, args = self.plugManager.openPage(plugin, **self.pluginArgs)
            if page:
                try:
                    return render_template(page, **args)
                except jinja2.exceptions.TemplateSyntaxError:
                    return render_template(
                        "error.html", status={"except": "plugin-page-corrupt"}
                    )
                except jinja2.exceptions.TemplateNotFound:
                    return render_template(
                        "error.html",
                        status={"except": "plugin-page-not-found", "page": page},
                    )
            else:
                abort(404)

    # /plugin/<plugin>/subpage/<page>
    def openPluginSubpage(self, plugin, page):
        if (
            self.plugManager.requiresAuth(plugin)
            and not current_user.is_authenticated()
        ):
            return render_template("requiresAuth.html")
        else:
            page, args = self.plugManager.openSubpage(plugin, page, **self.pluginArgs)
            if page:
                try:
                    return render_template(page, **args)
                except jinja2.exceptions.TemplateSyntaxError:
                    return render_template(
                        "error.html", status={"except": "plugin-page-corrupt"}
                    )
                except jinja2.exceptions.TemplateNotFound:
                    return render_template(
                        "error.html",
                        status={"except": "plugin-page-not-found", "page": page},
                    )
            else:
                abort(404)

    # /plugin/<plugin>/_cve_action/<action>
    def _jsonCVEAction(self, plugin, action):
        cve = request.args.get("cve", type=str)
        response = self.plugManager.onCVEAction(
            cve, plugin, action, fields=dict(request.args), **self.pluginArgs
        )
        if type(response) is bool and response is True:
            return jsonify({"status": "plugin_action_complete"})
        elif type(response) is bool and response is False or response is None:
            return jsonify({"status": "plugin_action_failed"})
        elif type(response) is dict:
            return jsonify(response)

    # /admin
    # /admin/
    def admin(self):
        if Configuration.loginRequired():
            if not current_user.is_authenticated():
                return render_template("login.html")
        else:
            person = User.get("_dummy_", self.auth_handler)
            login_user(person)
        output = None
        if os.path.isfile(Configuration.getUpdateLogFile()):
            with open(Configuration.getUpdateLogFile()) as updateFile:
                separator = "==========================\n"
                output = updateFile.read().split(separator)[-2:]
                output = separator + separator.join(output)
        return render_template("admin.html", status="default", **self.adminInfo(output))

    # /admin/change_pass
    @login_required
    def change_pass(self):
        current_pass = request.args.get("current_pass")
        new_pass = request.args.get("new_pass")
        if current_user.authenticate(current_pass):
            if new_pass:
                changePassword(current_user.id, new_pass)
                return jsonify({"status": "password_changed"})
            return jsonify({"status": "no_password"})
        else:
            return jsonify({"status": "wrong_user_pass"})

    # /admin/request_token
    @login_required
    def request_token(self):
        return jsonify({"token": generateToken(current_user.id)})

    # /admin/updatedb
    @login_required
    def updatedb(self):
        process = subprocess.Popen(
            [sys.executable, os.path.join(_runPath, "../sbin/db_updater.py"), "-civ"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        out, err = process.communicate()
        output = (
            "%s\n\nErrors:\n%s" % (str(out, "utf-8"), str(err, "utf-8"))
            if err
            else str(out, "utf-8")
        )
        return jsonify({"updateOutput": output, "status": "db_updated"})

    # /admin/whitelist
    # /admin/blacklist
    @login_required
    def listView(self):
        if request.url_rule.rule.split("/")[2].lower() == "whitelist":
            return render_template(
                "list.html", rules=getWhitelist(), listType="Whitelist"
            )
        else:
            return render_template(
                "list.html", rules=getBlacklist(), listType="Blacklist"
            )

    # /admin/whitelist/import
    # /admin/blacklist/import
    @login_required
    def listImport(self, force=None, path=None):
        _list = request.url_rule.split("/")[2]
        file = request.files["file"]
        force = request.form.get("force")
        count = countWhitelist() if _list.lower == "whitelist" else countBlacklist()
        if (count == 0) | (not count) | (force == "f"):
            if _list.lower == "whitelist":
                dropWhitelist()
                importWhitelist(TextIOWrapper(file.stream))
            else:
                dropBlacklist()
                importBlacklist(TextIOWrapper(file.stream))
            status = _list[0] + "l_imported"
        else:
            status = _list[0] + "l_already_filled"
        return render_template("admin.html", status=status, **self.adminInfo())

    # /admin/whitelist/export
    # /admin/blacklist/export
    @login_required
    def listExport(self, force=None, path=None):
        _list = request.url_rule.rule.split("/")[2]
        bytIO = BytesIO()
        data = exportWhitelist() if _list.lower == "whitelist" else exportBlacklist()
        bytIO.write(bytes(data, "utf-8"))
        bytIO.seek(0)
        return send_file(bytIO, as_attachment=True, attachment_filename=_list + ".txt")

    # /admin/whitelist/drop
    # /admin/blacklist/drop
    @login_required
    def listDrop(self):
        _list = request.url_rule.split("/")[2].lower()
        if _list == "whitelist":
            dropWhitelist()
        else:
            dropBlacklist()
        return jsonify({"status": _list[0] + "l_dropped"})

    # /admin/addToList
    @login_required
    def listAdd(self):
        cpe = request.args.get("cpe")
        cpeType = request.args.get("type")
        lst = request.args.get("list")
        if cpe and cpeType and lst:
            status = (
                "added_to_list"
                if self.addCPEToList(cpe, lst, cpeType)
                else "already_exists_in_list"
            )
            returnList = getWhitelist() if lst == "whitelist" else getBlacklist()
            return jsonify(
                {"status": status, "rules": returnList, "listType": lst.title()}
            )
        else:
            return jsonify({"status": "could_not_add_to_list"})

    # /admin/removeFromList
    @login_required
    def listRemove(self):
        cpe = request.args.get("cpe", type=str)
        cpe = urllib.parse.quote_plus(cpe).lower()
        cpe = cpe.replace("%3a", ":")
        cpe = cpe.replace("%2f", "/")
        lst = request.args.get("list", type=str)
        if cpe and lst:
            result = (
                removeWhitelist(cpe)
                if lst.lower() == "whitelist"
                else removeBlacklist(cpe)
            )
            status = (
                "removed_from_list" if (result > 0) else "already_removed_from_list"
            )
        else:
            status = "invalid_cpe"
        returnList = getWhitelist() if lst == "whitelist" else getBlacklist()
        return jsonify({"status": status, "rules": returnList, "listType": lst.title()})

    # /admin/editInList
    @login_required
    def listEdit(self):
        old = request.args.get("oldCPE")
        new = request.args.get("cpe")
        lst = request.args.get("list")
        CPEType = request.args.get("type")
        if old and new:
            result = (
                updateWhitelist(old, new, CPEType)
                if lst == "whitelist"
                else updateBlacklist(old, new, CPEType)
            )
            status = "cpelist_updated" if (result) else "cpelist_update_failed"
        else:
            status = "invalid_cpe"
        returnList = (
            list(getWhitelist()) if lst == "whitelist" else list(getBlacklist())
        )
        return jsonify({"rules": returnList, "status": status, "listType": lst})

    # /admin/listmanagement/<vendor>/<product>
    # /admin/listmanagement/<vendor>
    # /admin/listmanagement
    @login_required
    def listManagement(self, vendor=None, product=None):
        try:
            if product is None:
                # no product selected yet, so same function as /browse can be used
                if vendor:
                    vendor = urllib.parse.quote_plus(vendor).lower()
                browseList = getBrowseList(vendor)
                vendor = browseList["vendor"]
                product = browseList["product"]
                version = None
            else:
                # product selected, product versions required
                version = getVersionsOfProduct(urllib.parse.quote_plus(product).lower())
            return render_template(
                "listmanagement.html", vendor=vendor, product=product, version=version
            )
        except redisExceptions.ConnectionError:
            return render_template(
                "error.html",
                status={
                    "except": "redis-connection",
                    "info": {
                        "host": Configuration.getRedisHost(),
                        "port": Configuration.getRedisPort(),
                    },
                },
            )

    # /admin/listmanagement/add
    @login_required
    def listManagementAdd(self):
        # retrieve the separate item parts
        item = request.args.get("item", type=str)
        listType = request.args.get("list", type=str)

        pattern = re.compile("^[a-z:/0-9.~_%-]+$")

        if pattern.match(item):
            item = item.split(":")
            added = False
            if len(item) == 1:
                # only vendor, so a check on cpe type is needed
                if self.redisdb.sismember("t:/o", item[0]):
                    if self.addCPEToList("cpe:/o:" + item[0], listType):
                        added = True
                if self.redisdb.sismember("t:/a", item[0]):
                    if self.addCPEToList("cpe:/a:" + item[0], listType):
                        added = True
                if self.redisdb.sismember("t:/h", item[0]):
                    if self.addCPEToList("cpe:/h:" + item[0], listType):
                        added = True
            elif 4 > len(item) > 1:
                # cpe type can be found with a mongo regex query
                result = getCVEs(query={"cpe_2_2": {"$regex": item[1]}})["results"]
                if result.count() != 0:
                    prefix = ((result[0])["cpe_2_2"])[:7]
                    if len(item) == 2:
                        if self.addCPEToList(
                            prefix + item[0] + ":" + item[1], listType
                        ):
                            added = True
                    if len(item) == 3:
                        if self.addCPEToList(
                            prefix + item[0] + ":" + item[1] + ":" + item[2], listType
                        ):
                            added = True
            status = "added_to_list" if added else "could_not_add_to_list"
        else:
            status = "invalid_cpe"
        j = {"status": status, "listType": listType}
        return jsonify(j)

    # /login
    def login_check(self):
        # validate username and password
        username = request.form.get("username")
        password = request.form.get("password")
        person = User.get(username, self.auth_handler)
        try:
            if person and person.authenticate(password):
                login_user(person)
                return render_template(
                    "admin.html", status="logged_in", **self.adminInfo()
                )
            else:
                return render_template("login.html", status="wrong_user_pass")
        except Exception as e:
            print(e)
            return render_template("login.html", status="outdated_database")

    # /logout
    @login_required
    def logout(self):
        logout_user()
        return redirect("/")


if __name__ == "__main__":
    server = Index()
    server.start()
